from PIL import Image
import random
import numpy as np
        # 0.000, 0.447, 0.741,
        # 0.850, 0.325, 0.098,
        # 0.929, 0.694, 0.125,
        # 0.494, 0.184, 0.556,
        # 0.466, 0.674, 0.188,
        # 0.301, 0.745, 0.933,
        # 0.635, 0.078, 0.184,
# fmt: off
# RGB:
_COLORS = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.300, 0.300, 0.300,
        0.600, 0.600, 0.600,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.749, 0.749, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 1.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.857, 0.857, 0.857,
        1.000, 1.000, 1.000
    ]
).astype(np.float32).reshape(-1, 3)

def random_colors(N, rgb=True, maximum=255):
    """
    Args:
        N (int): number of unique colors needed
        rgb (bool): whether to return RGB colors or BGR colors.
        maximum (int): either 255 or 1

    Returns:
        ndarray: a list of random_color
    """
    indices = random.sample(range(len(_COLORS)), N)
    ret = [_COLORS[i] * maximum for i in indices]
    if not rgb:
        ret = [x[::-1] for x in ret]
    return ret

def draw_mask(src,mask,id,new_color):
    m2 = mask == id
    src[m2] = new_color
    

def convert(in_img_file, in_mask_file, out_mask_file):
    in_mask = Image.open(in_mask_file)
    in_mask = np.array(in_mask)
    in_mask[...,-1] = 0
    nn_mask = np.where(in_mask[:,:,0] > 0,0.6,0.0)
    idsmask = in_mask[..., 0]
    ids = np.unique(idsmask)
    colos = random_colors(len(ids))
    tempc = np.zeros(4,np.float32)
    for thing_id,color in zip(ids,colos):
        if thing_id == 0:
            continue
        tempc[0:3] = color
        draw_mask(in_mask,idsmask, thing_id, tempc)
    #fusion
    
    

    in_image = Image.open(in_img_file)
    in_image = np.array(in_image)
    image2 = np.einsum('hwc,hw->hwc', in_image, (1-nn_mask))
    mask2 = np.einsum('hwc,hw->hwc', in_mask, nn_mask)
    image2 = image2[...,0:3] + mask2[...,0:3]
    image2 = image2.astype(np.uint8)
    Image.fromarray(image2).save(out_mask_file)


def convert_2(in_img_file, in_mask_file, out_mask_file,
              in_img_file_2, in_mask_file_2, out_mask_file_2
              ):
    in_mask = Image.open(in_mask_file)
    in_mask = np.array(in_mask)
    in_mask[...,-1] = 0
    nn_mask = np.where(in_mask[:,:,0] > 0,0.6,0.0)
    idsmask = in_mask[..., 0]
    ids = np.unique(idsmask)
    
    in_mask_2 = Image.open(in_mask_file_2)
    in_mask_2 = np.array(in_mask_2)
    in_mask_2[...,-1] = 0
    nn_mask_2 = np.where(in_mask_2[:,:,0] > 0,0.6,0.0)
    idsmask_2 = in_mask_2[..., 0]
    ids_2 = np.unique(idsmask_2)
    
    ids_a = np.union1d(ids, ids_2)
    
    colos = random_colors(len(ids_a))
    tempc = np.zeros(4,np.float32)
    for thing_id,color in zip(ids,colos):
        if thing_id == 0:
            continue
        tempc[0:3] = color
        if thing_id in ids:
            draw_mask(in_mask,idsmask, thing_id, tempc)
        if thing_id in ids_2:
            draw_mask(in_mask_2,idsmask_2, thing_id, tempc)
    #fusion
    
    

    in_image = Image.open(in_img_file)
    in_image = np.array(in_image)
    image2 = np.einsum('hwc,hw->hwc', in_image, (1-nn_mask))
    mask2 = np.einsum('hwc,hw->hwc', in_mask, nn_mask)
    image2 = image2[...,0:3] + mask2[...,0:3]
    image2 = image2.astype(np.uint8)
    Image.fromarray(image2).save(out_mask_file)

    in_image_2 = Image.open(in_img_file_2)
    in_image_2 = np.array(in_image_2)
    image2_2 = np.einsum('hwc,hw->hwc', in_image_2, (1-nn_mask_2))
    mask2_2 = np.einsum('hwc,hw->hwc', in_mask_2, nn_mask_2)
    image2_2 = image2_2[...,0:3] + mask2_2[...,0:3]
    image2_2 = image2_2.astype(np.uint8)
    Image.fromarray(image2_2).save(out_mask_file_2)
    
def convert_3(in_mask_file, out_mask_file):
    in_mask = Image.open(in_mask_file)
    in_mask = np.array(in_mask)
    in_mask[...,-1] = 0
    nn_mask = np.where(in_mask[:,:,0] > 0,0.6,0.0)
    idsmask = in_mask[..., 0]
    ids = np.unique(idsmask)
    colos = random_colors(len(ids))
    tempc = np.zeros(4,np.float32)
    for thing_id,color in zip(ids,colos):
        if thing_id == 0:
            continue
        tempc[0:3] = color
        tempc[3] = 255
        draw_mask(in_mask,idsmask, thing_id, tempc)
    #fusion
    Image.fromarray(in_mask).save(out_mask_file)
    
    
# convert("E:/data_os/BEAR/QGY/IMG1-2.png", "E:/data_os/BEAR/QGY/MASK1.png", "E:/data_os/BEAR/QGY/MASK1-colors-c.png")
# convert("E:/data_os/BEAR/QGY/IMG2-2.png", "E:/data_os/BEAR/QGY/MASK2.png", "E:/data_os/BEAR/QGY/MASK2-colors.png")

# convert_2("E:/data_os/BEAR/QGY/IMG1-2.png", "E:/data_os/BEAR/QGY/MASK1.png", "E:/data_os/BEAR/QGY/MASK1-colors-2.png",
#           "E:/data_os/BEAR/QGY/IMG2-2.png", "E:/data_os/BEAR/QGY/MASK2.png", "E:/data_os/BEAR/QGY/MASK2-colors-2.png")

convert_3("F:/study/XLWS/D5/图/QGY/MASK1.png", "F:/study/XLWS/D5/图/QGY/MASK1-colors-c.png")
